# -*- coding: utf-8 -*-
import os
import argparse
import sys
import cv2 as cv
import numpy as np
import platform
import struct
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.ticker import NullLocator
import random


# Bounding-box colors
cmap = plt.get_cmap("tab20b")
colors = [cmap(i) for i in np.linspace(0, 1, 20)]

def drew_detection_on_img(detections, img_path, idx_2_cls, save_path):
    """
    drew bboxes on image
    :param detections:  x1, y1, x2, y2, cls_conf, cls_pred   (bottom left coord)
    :param img_path:    origin img path,
    :param idx_2_cls:   idx_2_cls dict
    :return:
    """
    # Create plot
    img = np.array(Image.open(img_path))
    plt.figure()
    fig, ax = plt.subplots(1)
    ax.imshow(img)

    # Draw bounding boxes and labels of detections
    if len(detections) != 0:
        # Rescale boxes to original image
        # detections = rescale_boxes(detections, opt.img_size, img.shape[:2])
        unique_labels = np.unique(detections[:, -1])
        n_cls_preds = len(unique_labels)
        bbox_colors = random.sample(colors, n_cls_preds)
        for x1, y1, box_w, box_h, cls_conf, cls_pred in detections:
            print("\t+ Label: %s, Conf: %.5f" % (idx_2_cls[int(cls_pred)], cls_conf.item()))

            color = bbox_colors[int(np.where(unique_labels == int(cls_pred))[0])]
            # Create a Rectangle patch
            bbox = patches.Rectangle((x1, y1), box_w, box_h, linewidth=2, edgecolor=color, facecolor="none")
            # Add the bbox to the plot
            ax.add_patch(bbox)
            # Add label
            plt.text(
                x1,
                y1,
                s=idx_2_cls[int(cls_pred)],
                color="white",
                verticalalignment="top",
                bbox={"color": color, "pad": 0},
            )

    # Save generated image with detections
    plt.axis("off")
    plt.gca().xaxis.set_major_locator(NullLocator())
    plt.gca().yaxis.set_major_locator(NullLocator())
    filename = img_path.split("/")[-1].split(".")[0]
    plt.savefig("{}/{}.png".format(save_path, filename), bbox_inches="tight", pad_inches=0.0)
    plt.cla()
    plt.close('all')

def putText(srcFile, dstFile, text) :
    image = cv.imread(srcFile)
    cv.putText(image, text, (15,20), cv.FONT_HERSHEY_COMPLEX_SMALL, 1.0, (0, 0, 255) )
    cv.imwrite(dstFile, image)
   

def showpic(data):
    cv.imshow("result", data)
    cv.waitKey(0)
    cv.destroyAllWindows()


def check_path(path):
    key = '\\'
    if platform.system() == "Linux":
        key = '/'
    if path[len(path)-1] != key:
        path += key
    return path


def Parse():
    """Parse input arguments."""
    parser = argparse.ArgumentParser(
        description='Convert Jpeg To NV12 and BGR raw data')
         
    parser.add_argument('--src  ', dest='src_name',
                        help='cmv or ped', type=str)
    parser.add_argument('--resize_h  ', dest='resize_h',
                        help='resize_h', type=str)
    parser.add_argument('--resize_w  ', dest='resize_w',
                        help='resize_w', type=str)
                        
    if len(sys.argv) == 1:
        helpInfo()
        parser.print_help()
        sys.exit(1)
    args = parser.parse_args()
    return args


def is_img(ext):
    ext = os.path.splitext(ext.lower())[1]
    if ext == '.jpg':
        return True
    elif ext == '.png':
        return True
    elif ext == '.jpeg':
        return True
    elif ext == '.bmp':
        return True
    else:
        return False


def mergeUV(u, v):
    if u.shape == v.shape:
        uv = np.zeros(shape=(u.shape[0], u.shape[1]*2))
        for i in range(0, u.shape[0]):
            for j in range(0, u.shape[1]):
                uv[i, 2*j] = u[i, j]
                uv[i, 2*j+1] = v[i, j]
        return uv
    else:
        print("size of Channel U is different with Channel V")


def rgb2nv12(image):
    if image.ndim == 3:
        b = image[:, :, 0]
        g = image[:, :, 1]
        r = image[:, :, 2]
        y = (0.299*r+0.587*g+0.114*b)
        u = (-0.169*r-0.331*g+0.5*b+128)[::2, ::2]
        v = (0.5*r-0.419*g-0.081*b+128)[::2, ::2]
        uv = mergeUV(u, v)
        yuv = np.vstack((y, uv))
        return yuv.astype(np.uint8)
    else:
        print("image is not BGR format")


def package2planar(image):
    if image.ndim == 3:
        b, g, r = cv.split(image)
        bgr = np.stack((b, g, r))
        bgr = np.reshape(bgr, (image.shape[2], image.shape[0], image.shape[1]))
        return bgr
    else:
        print("image is not BGR format")


def compare(src, dest):
    if src.shape == dest.shape:
        if (src == dest).all():
            return True
        else:
            return False
    else:
        print("These two images are not in same size")
        return False


def mkdirown(path):
    if os.path.exists(path) == False:
        os.makedirs(path)


def jpeg2yuv(src_name,resize_w, resize_h):
    src_root_path =  '../jpg'   #'./data/jpg'
    nv12_root_path = '../nv12'  #'./data/nv12'
    bgr_dest_path =  '../bgr'  #'./data/bgr'
    
    image_ori = cv.imread(src_name)         
    image_ori = cv.resize(image_ori,(resize_w, resize_h))   
    yuv = rgb2nv12(image_ori)
    
    return yuv


def camera_jpeg2yuv(image_ori, resize_w, resize_h):
    # src_root_path = '../jpg'  # './data/jpg'
    # nv12_root_path = '../nv12'  # './data/nv12'
    # bgr_dest_path = '../bgr'  # './data/bgr'

    # image_ori = cv.imread(src_name)
    image_ori = cv.resize(image_ori, (resize_w, resize_h))
    yuv = rgb2nv12(image_ori)

    return yuv


def saveFile(yuv, fileName) :
    with open(fileName, "wb") as fp:
        fp.write(yuv)
        fp.close()       


def helpInfo():
    print("Image Conversion Tool")
    print("Only support bmp/jpeg/jpg/png Format")
    print("Programmed by \033[1;31mz00418008\033[0m\n")
   
    

if __name__ == '__main__':
    args = Parse()